from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments, AutoTokenizer
from datasets import load_from_disk
import torch
import torch.distributed as dist
import os
from peft import PeftModel
import evaluate
import numpy as np 
# import bitsandbytes as bnb  # QLoRA需要的库



"""
运行命令
CUDA_VISIBLE_DEVICES=1,2 deepspeed --num_gpus=2 /changzheng/PracticalTraining/lyj/lingdongai/server/peft/QLora_qwen2_5_coder_7b.py
CUDA_VISIBLE_DEVICES=0,1,2 deepspeed --num_gpus=3 /changzheng/PracticalTraining/lyj/lingdongai/server/peft/QLora_qwen2_5_coder_7b.py
"""

"""
可视化 TensorBoard
tensorboard --logdir /changzheng/PracticalTraining/lyj/lingdongai/models/logs
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # 禁用tokenizers并行化

# 手动指定要使用的 GPU
# os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"

# 确保 PyTorch 使用的是第一个可见的 GPU
torch.cuda.set_device(0)  # 对应的是 CUDA_VISIBLE_DEVICES 中的第一个设备，即物理 GPU 1

# 减少内存碎片化
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
# 释放未使用的显存
torch.cuda.empty_cache()

# 初始化进程组，用于分布式训练的通信
dist.init_process_group(backend="nccl")

model_name_or_path = '/changzheng/PracticalTraining/lyj/lingdongai/models/model/Qwen/Qwen2___5-7B-Instruct'
data_path = '/changzheng/PracticalTraining/lyj/lingdongai/models/data/medical_dialogue_53962'

# 初始化分词器，用于后续数据预处理
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
print(tokenizer)

# 加载数据集
train_data = load_from_disk(data_path)
print(train_data)

# 定义数据预处理函数，仅使用 input 作为输入
def preprocess_function(examples):
    # 对输入的 input 进行分词并转换为模型的输入格式
    model_inputs = tokenizer(examples["input"], max_length=256, truncation=True, padding="max_length")
    
    # 对输出的 output 进行分词并作为标签
    labels = tokenizer(examples["output"], max_length=256, truncation=True, padding="max_length")["input_ids"]
    model_inputs["labels"] = labels
    
    return model_inputs

# 查看下标为0的数据
print("查看下标0的输入", train_data['input'][0])
print("查看下标0的标签", train_data['output'][0])

# 对数据集进行预处理
tokenized_datasets = train_data.map(preprocess_function, batched=True)

# 将训练集划分为训练集和验证集
train_test_split = tokenized_datasets.train_test_split(test_size=0.3)

train_dataset = train_test_split['train']  # 训练集
eval_dataset = train_test_split['test']  # 验证集

# 定义LoRA配置，设置用于微调的参数
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,  # 因果语言模型任务
    r=16,  # LoRA的秩
    lora_alpha=32,  # 缩放因子
    lora_dropout=0.2,  # LoRA的dropout率
    target_modules=['q_proj', 'k_proj', 'v_proj'],  # 目标模块，指定需要微调的部分
    bias="none"  # 是否应用bias，这里不使用
)

# 加载预训练的Qwen-7B模型
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)

# 加载Qwen-7B模型并启用4-bit量化
# model = AutoModelForCausalLM.from_pretrained(
#     model_name_or_path,
#     load_in_4bit=True,  # 启用4-bit量化
#     device_map="auto",  # 自动映射设备，节省内存
#     torch_dtype=torch.float16,  # 使用float16减少显存占用
#     quantization_config=bnb.QuantizationConfig(llm_int8_threshold=6.0)  # 配置量化设置
# )

# 冻结模型的所有参数，避免它们在训练中被更新
for param in model.parameters():
    param.requires_grad = False

# 使用LoRA配置微调模型
peft_model = get_peft_model(model, lora_config)

# 加载准确率评估指标
# accuracy_metric = evaluate.load("/changzheng/PracticalTraining/lyj/lingdongai/models/metrics/accuracy")

# 加载 F1 分数评估指标
f1_metric = evaluate.load("/changzheng/PracticalTraining/lyj/lingdongai/models/metrics/f1")

# 修改计算评估指标的函数，加入 F1 计算
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    
    # 确保 logits 是 Tensor 类型，如果它是 numpy 数组，需将其转换为 Tensor
    if isinstance(logits, np.ndarray):
        logits = torch.from_numpy(logits)
    
    # 计算每个 token 的预测值
    predictions = torch.argmax(logits, dim=-1)
    
    # 确保 labels 也是 Tensor 类型
    if isinstance(labels, np.ndarray):
        labels = torch.from_numpy(labels)
    
    # 去除 padding token，以免影响准确率计算 (假设 padding token ID 是 151643)
    mask = labels != 151643
    predictions = predictions[mask]
    labels = labels[mask]
    
    # 计算准确率
    # accuracy = accuracy_metric.compute(predictions=predictions, references=labels)
    
    # 计算 F1 分数
    f1_score = f1_metric.compute(predictions=predictions, references=labels, average="weighted")
    
    return {"f1": f1_score["f1"]}




# 定义训练参数，并加入DeepSpeed配置
training_args = TrainingArguments(
    output_dir='/changzheng/PracticalTraining/lyj/lingdongai/models/logs_model/qwen_7b_logs_model',  # 输出目录，用于保存模型
    logging_dir='/changzheng/PracticalTraining/lyj/lingdongai/models/logs/qwen_7b_logs',  # 日志目录
    overwrite_output_dir=True,  # 允许覆盖输出目录中的内容
    num_train_epochs=50,  # 训练的轮数
    per_device_train_batch_size=1,  # 每个设备的批次大小，这个是数据并行的部分
    gradient_accumulation_steps=128,  # 梯度累积步数，用于模拟更大的批次
    save_steps=100,  # 每XX步保存一次模型
    logging_steps=100,  # 每100步记录一次损失日志
    learning_rate=5e-4,  # 初始学习率
    fp16=True,  # 启用混合精度训练，以减少显存使用
    evaluation_strategy="steps",  # 设置评估策略，每隔一定步数评估一次
    eval_steps=100,  # 每xx步进行一次评估
    save_total_limit=2,  # 最多保留2个模型检查点
    remove_unused_columns=False,  # 不删除未使用的列
    weight_decay=0.01,  # 权重衰减系数，防止过拟合
    warmup_steps=100,  # 学习率预热的步数
    deepspeed="/changzheng/PracticalTraining/lyj/lingdongai/server/peft/deepspeed_config.json",  # DeepSpeed配置文件
)

# 定义Trainer，用于训练和评估
trainer = Trainer(
    model=peft_model,  # 使用微调后的LoRA模型
    args=training_args,  # 训练参数
    train_dataset=train_dataset,  # 训练集
    eval_dataset=eval_dataset,  # 验证集
    compute_metrics=compute_metrics  # 添加计算准确率的评估函数
)

# 开始微调训练
trainer.train()

# 保存微调后的模型
peft_model.save_pretrained('/changzheng/PracticalTraining/lyj/lingdongai/models/new_model/lora_qwen_7b_model')

# 手动评估验证集
eval_results = trainer.evaluate()
print("评估结果:", eval_results)

print("************************")
print("lora训练运行完毕")
print("************************")

# merge_and_unlaod  # 融合多个lora模型